
import torch
from torch import nn
from model import makeResNet152, makeResNet50, makeNoClassiferResNet152, makeNoClassiferResNet50, makeResNext101_32x8d
from torch import optim
from torch.utils.data import DataLoader
from dataset import classiferDataset, VeRiClassiferDataset
from loss import FocalLoss, CircleLoss, convert_label_to_similarity

def main():
    device = torch.device("cuda")
    maxAcc = 0
    maxLoss = 9999999999
    batchsz = 32
    classNum = 576
    # dataPath = "G:\\DeepLearning\\DeepLearningDataset\\VERI-Wild\\data\\"
    #
    # trainDataSet = classiferDataset(dataPath + "train\\imgs\\", dataPath + "train\\labels\\")
    # print(len(trainDataSet))
    # trainDataSet = DataLoader(trainDataSet, batch_size=batchsz, shuffle=True, num_workers=4)
    #
    # testDataSet = classiferDataset(dataPath + "test\\imgs\\", dataPath + "test\\labels\\")
    # print(len(testDataSet))
    # testDataSet = DataLoader(testDataSet, batch_size=batchsz*8, shuffle=True, num_workers=4)

    dataPath = ".\\data\\VeRi\\image_train\\"

    trainDataSet = VeRiClassiferDataset(dataPath)
    print(len(trainDataSet))
    trainDataSet = DataLoader(trainDataSet, batch_size=batchsz, shuffle=True, num_workers=4)

    testDataSet = VeRiClassiferDataset(dataPath)
    print(len(testDataSet))
    testDataSet = DataLoader(testDataSet, batch_size=batchsz*8, shuffle=True, num_workers=4)

    # model = makeResNext101_32x8d(classNum).to(device)
    # model = makeNoClassiferResNet152().to(device)
    # model = makeNoClassiferResNet152().to(device)
    model = torch.load(".\\savedModels\\model_epoch_46_right_0.9995500026470433.pkl")
    print(model)
    p = sum(map(lambda p:p.numel(), model.parameters()))
    print('parameters size:', p)

    startEpoch = 50

    _lr = 0
    optimizer = optim.Adam(model.parameters(), lr=_lr , betas=(0.9, 0.999), eps=1e-09)
    criterion = nn.CrossEntropyLoss().to(device)
    # criterion = FocalLoss(classNum).to(device)
    # criterion = CircleLoss(m=0.25, gamma=256).to(device)

    for epoch in range(startEpoch, 10000):

        if epoch >= 0:
            _lr = 1e-4
        if epoch >= 15:
            _lr = 1e-5
        if epoch >= 50:
            _lr = 1e-6
        if epoch >= 100:
            _lr = 1e-7
        for p in optimizer.param_groups:
            p['lr'] = _lr


        trainTotalLoss = 0
        trainTotalStep = 0
        model.train()
        for batchidx, (x, label) in enumerate(trainDataSet):

            x, label = x.to(device), label.to(device)
            logits = model(x)
            label = label.long()

            loss = criterion(logits, label)
            # inp_sp, inp_sn = convert_label_to_similarity(logits, label)
            # loss = criterion(inp_sp, inp_sn)
            # loss = criterion(logits, label)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            trainTotalLoss += float(loss)
            trainTotalStep += 1

            if (batchidx % 100 == 0):
                print(float(loss))
                # print(trainTotalLoss / trainTotalStep)


        model.eval()
        testTotalRight = 0
        testTotalNum = 0
        for batchidx, (x, label) in enumerate(testDataSet):
            with torch.no_grad():
                x, label = x.to(device), label.to(device)
                logits = model(x)

                # logits = torch.softmax(logits, dim=-1)
                logits = torch.argmax(logits, dim=-1)

                same = (logits == label).long().sum()
                testTotalRight += same.item()
                testTotalNum += logits.shape[0]
        testAcc = testTotalRight / testTotalNum


        print("epoch: ", epoch, "  testAcc: ", testAcc)
        if epoch > 0 and testAcc >= maxAcc:
            torch.save(model, ".\\savedModels\\model_epoch_{}_right_{}.pkl".format(epoch, testAcc))
            maxAcc = testAcc
        if epoch % 1 == 0:
            torch.save(model, ".\\savedModels\\n_model_epoch_{}_right_{}.pkl".format(epoch, testAcc))


        # trainLoss = trainTotalLoss / trainTotalStep
        # print("epoch: ", epoch, "   train loss: ", trainLoss)
        # if epoch > 0 and  epochLoss < maxLoss:
        #     torch.save(model, ".\\savedModels\\model_epoch_{}_right_{}.pkl".format(epoch, epochLoss))
        #     maxLoss = epochLoss
        # if epoch % 1 == 0:
        #     torch.save(model, ".\\savedModels\\n_model_epoch_{}_right_{}.pkl".format(epoch, epochLoss))



if __name__ == '__main__':
    main()
